[externalization] Add a two-step patch/sub-export API for module externalization - #53
Open
gokulkrishna98 wants to merge 3 commits into
Open
Conversation
Port the source-side changes of the module externalization API into coreai_torch/, replacing the marker-based approach. Tests cover: - backward: gradients flow through externalized submodules via register_autograd. - no call sites: exporting before patching leaves no custom-op call sites, so _subexport_and_restore warns and skips externalization instead of raising, and conversion still succeeds. - the manual patch/sub-export/convert workflow, i.e. the documented advanced path where the caller drives export themselves instead of using add_pytorch_module. A passthrough export stands in for the quantizer, pinning the workflow contract: custom-op call sites survive the caller's export, the model is left unpatched after _subexport_and_restore, and the marked submodule is emitted as a rms_norm composite op. - numerics for that workflow: compiles the converted program, runs it in the Core AI runtime, and compares against the restored PyTorch model.
Cross-reference the two "no call sites" warning paths so neither reads as dead code left over from the refactor: _drop_missing_call_sites handles top-level marked submodules missing from the whole-model program, while _PreparedModules.__iter__ handles the nested case it cannot see, once the enclosing module's sub-export exists. Both docstrings now point at each other, and _subexport_and_restore names the split. Add test_externalize_multiple_staged_entries_numerics: two staged entries in one to_coreai() call, one via add_pytorch_module(externalize_modules=) and one plain add_exported_program, validating numerics on both entrypoints. This pins the invariant that externalization is per-entry — _init_conversion_state() resets _externalized_exported_programs before each entry, so a plain entry converted alongside an externalized one stays flat. Verified non-vacuous: dropping that reset and adding a third _perform_externalization call site fails the test on a duplicate rms_norm symbol. _validate_numerics takes a function_name argument so it can check a named entrypoint instead of only "main".
The module docstring documents four Phases (Mark, Prepare, Export, Emit) while individual docstrings still used a separate six-Step scheme that no longer lined up — export was "Step 4" but Phase 3, prepare was "Step 3" but Phase 2, and there was no Step 2 at all. Relabel every per-function reference to the Phase vocabulary and drop the number from _restore_externalized, since restore is no longer its own phase. Docstrings only; no behavior change.
gokulkrishna98
marked this pull request as ready for review
August 4, 2026 16:33
gokulkrishna98
requested review from
Lewis300,
jakesabathia2,
pkmandke and
vineet-g
August 4, 2026 16:34
jakesabathia2
approved these changes
Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Reworks module externalization — emitting an
nn.Modulesubtree as a single Core AI composite op — into an explicit two-step API, and adds autograd support so gradients can flow through externalized submodules.The one-shot path is unchanged:
add_pytorch_module(model, export_fn=…, externalize_modules=[...])still does everything. The new split exists for callers who must drive export themselves (e.g. a quantizer), so marking and sub-export can happen either side of their own passes:Changes
into _subexport_and_restore's finally, so the model is always left unpatched.